from os import listdir
from os.path import splitext
from pathlib import Path
import os

import numpy as np
import torch
from PIL import Image
from scipy.io import loadmat
from torch.utils.data import Dataset


class BasicDataset(Dataset):
    def __init__(self, root_dir: str, in_channel: list, out_channel: list, suffix: str):
        self.root_dir = Path(root_dir)
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.suffix = suffix
        self.img_dir = os.path.join(self.root_dir, "模糊图像")
        self.label_dir = os.path.join(self.root_dir, "光谱图像")
        self.ids = [splitext(file)[0] for file in listdir(self.img_dir) if not file.startswith('.')]

        # 加载a光谱幅照度修正参数
        a_full = loadmat(os.path.join(self.root_dir, 'a_326.mat'))['a']
        sample_dir = os.path.join(self.label_dir, listdir(self.label_dir)[0])
        wavelength_full = sorted([int(i[0:-len(self.suffix)])
                                  for i in [splitext(file)[0]
                                            for file in listdir(sample_dir) if not file.startswith('.')]])
        self.a = np.asarray([a_full[wavelength_full.index(i)] for i in out_channel])

        # max_radiance
        self.max_radiance = np.asarray(loadmat(os.path.join(self.root_dir, 'a_326.mat'))['maxval'])

        # 后缀名
        ext_full = list(set(i
                            for i in [splitext(file)[1]
                                      for file in listdir(sample_dir) if not file.startswith('.')]))
        self.ext = ext_full[0]

        # 给ids按编号大小排序
        self.ids = [int(i) for i in self.ids]
        self.ids.sort()
        self.ids = [str(i) for i in self.ids]

    def __len__(self):
        return len(self.ids)

    @classmethod
    def preprocess(cls, pil_img_list, is_label, a):
        c = len(pil_img_list)
        w, h = pil_img_list[0].size
        img_ndarray = np.zeros((c, h, w))
        if not is_label:
            for i in range(c):
                ndarray = np.asarray(pil_img_list[i])
                img_ndarray[i, :, :] = ndarray / 65535
        elif is_label:
            for i in range(c):
                ndarray = np.asarray(pil_img_list[i])
                img_ndarray[i, :, :] = ndarray / 65535 * a[i]
        return img_ndarray

    @classmethod
    def load(cls, filename):
        ext = splitext(filename)[1]
        if ext in ['.npz', '.npy']:
            return Image.fromarray(np.load(filename))
        elif ext in ['.pt', '.pth']:
            return Image.fromarray(torch.load(filename).numpy())
        else:
            return Image.open(filename)

    def __getitem__(self, idx):
        name = self.ids[idx]
        # mask_file = list(self.masks_dir.glob(name + self.mask_suffix + '.*'))
        img_files = [os.path.join(self.img_dir, name, str(i) + self.suffix + self.ext)
                     for i in self.in_channel]
        label_files = [os.path.join(self.label_dir, name, str(i) + self.suffix + self.ext)
                       for i in self.out_channel]

        # assert len(mask_file) == 1, f'Either no mask or multiple masks found for the ID {name}: {mask_file}'
        # assert len(img_file) == 1, f'Either no image or multiple images found for the ID {name}: {img_file}'
        img = []
        label = []
        for i in range(len(img_files)):
            img.append(self.load(img_files[i]))
        for i in range(len(label_files)):
            label.append(self.load(label_files[i]))

        img = self.preprocess(img, False, self.a)
        label = self.preprocess(label, True, self.a)

        return {
            'image': torch.as_tensor(img.copy()).double().contiguous(),
            'label': torch.as_tensor(label.copy()).double().contiguous()
        }

    def print_info(self):
        img_and_label = self[0]
        img = img_and_label['image']
        label = img_and_label['label']
        print('-----------------------------------------------------------------')
        print('模糊图像Channel,Height,Width: {}'.format(img.shape))
        print('模糊图像dtype: {}'.format(img.dtype))
        print('模糊图像值范围: {} ～ {}'.format(img.min(), img.max()))
        print('真实光谱图像Channel,Height,Width: {}'.format(label.shape))
        print('真实光谱图像dtype: {}'.format(label.dtype))
        print('真实光谱图像值范围: {} ～ {}'.format(label.min(), label.max()))
        print('-----------------------------------------------------------------')

    def get_max_radiance(self):
        return self.max_radiance

class WuZeiDataset(BasicDataset):
    def __init__(self, root_dir, in_channel, out_channel, suffix='nm'):
        super().__init__(root_dir, in_channel, out_channel, suffix)
